#!/bin/bash

# uWSGI Exporter 测试脚本
# 用于验证指标收集功能

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 配置
EXPORTER_PORT=9070
UWSGI_PORT=8080
UWSGI_SOCKET="/tmp/uwsgi_stats.sock"
TEST_APP_URL="http://localhost:${UWSGI_PORT}"
EXPORTER_URL="http://localhost:${EXPORTER_PORT}"

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查依赖
check_dependencies() {
    log_info "检查依赖..."
    
    # 检查 Python
    if ! command -v python3 &> /dev/null; then
        log_error "Python3 未安装"
        exit 1
    fi
    
    # 检查 pip
    if ! command -v pip3 &> /dev/null; then
        log_error "pip3 未安装"
        exit 1
    fi
    
    # 检查 uwsgi
    if ! command -v uwsgi &> /dev/null; then
        log_warning "uWSGI 未安装，尝试安装..."
        pip3 install uwsgi
    fi
    
    # 检查 netcat
    if ! command -v nc &> /dev/null; then
        log_warning "netcat 未安装，某些测试可能失败"
    fi
    
    log_success "依赖检查完成"
}

# 安装 Python 依赖
install_python_deps() {
    log_info "安装 Python 依赖..."
    
    # 安装 Flask
    pip3 install flask
    
    log_success "Python 依赖安装完成"
}

# 启动 uWSGI 测试应用
start_uwsgi() {
    log_info "启动 uWSGI 测试应用..."
    
    # 检查是否已经在运行
    if pgrep -f "uwsgi.*uwsgi.ini" > /dev/null; then
        log_warning "uWSGI 测试应用已在运行"
        return 0
    fi
    
    # 启动 uWSGI
    cd test
    uwsgi --ini uwsgi.ini --daemonize /tmp/uwsgi_daemon.log
    
    # 等待启动
    sleep 3
    
    # 检查是否启动成功
    if ! pgrep -f "uwsgi.*uwsgi.ini" > /dev/null; then
        log_error "uWSGI 启动失败"
        cat /tmp/uwsgi_daemon.log
        exit 1
    fi
    
    log_success "uWSGI 测试应用启动成功"
}

# 测试 uWSGI 应用
test_uwsgi_app() {
    log_info "测试 uWSGI 应用..."
    
    # 等待应用完全启动
    sleep 2
    
    # 测试主页
    if curl -s "${TEST_APP_URL}/" | grep -q "uWSGI Test Application"; then
        log_success "uWSGI 应用主页访问成功"
    else
        log_error "uWSGI 应用主页访问失败"
        return 1
    fi
    
    # 测试健康检查
    if curl -s "${TEST_APP_URL}/api/health" | grep -q "healthy"; then
        log_success "uWSGI 应用健康检查成功"
    else
        log_error "uWSGI 应用健康检查失败"
        return 1
    fi
    
    # 生成一些负载
    log_info "生成测试负载..."
    for i in {1..10}; do
        curl -s "${TEST_APP_URL}/" > /dev/null
        curl -s "${TEST_APP_URL}/api/load" > /dev/null
        if [ $((i % 3)) -eq 0 ]; then
            curl -s "${TEST_APP_URL}/api/error" > /dev/null 2>&1 || true
        fi
        sleep 0.1
    done
    
    log_success "uWSGI 应用测试完成"
}

# 测试统计接口
test_stats_interface() {
    log_info "测试 uWSGI 统计接口..."
    
    # 测试 Unix socket 统计接口
    if command -v nc &> /dev/null; then
        if echo "stats" | nc -U "${UWSGI_SOCKET}" 2>/dev/null | grep -q "workers"; then
            log_success "Unix socket 统计接口访问成功"
        else
            log_error "Unix socket 统计接口访问失败"
            return 1
        fi
    else
        log_warning "跳过 Unix socket 测试（netcat 未安装）"
    fi
    
    # 测试 HTTP 统计接口
    if curl -s "${TEST_APP_URL}/api/stats" | grep -q "request_count"; then
        log_success "HTTP 统计接口访问成功"
    else
        log_error "HTTP 统计接口访问失败"
        return 1
    fi
    
    log_success "统计接口测试完成"
}

# 启动 Exporter
start_exporter() {
    log_info "启动 uWSGI Exporter..."
    
    # 检查是否已经在运行
    if pgrep -f "uos_uwsgi_exporter" > /dev/null; then
        log_warning "uWSGI Exporter 已在运行"
        return 0
    fi
    
    # 启动 Exporter
    cd ..
    ./uos_uwsgi_exporter -sockets "${UWSGI_SOCKET}" -port ${EXPORTER_PORT} &
    EXPORTER_PID=$!
    
    # 等待启动
    sleep 3
    
    # 检查是否启动成功
    if ! kill -0 ${EXPORTER_PID} 2>/dev/null; then
        log_error "uWSGI Exporter 启动失败"
        exit 1
    fi
    
    log_success "uWSGI Exporter 启动成功 (PID: ${EXPORTER_PID})"
}

# 测试 Exporter
test_exporter() {
    log_info "测试 uWSGI Exporter..."
    
    # 等待 Exporter 完全启动
    sleep 2
    
    # 测试健康检查
    if curl -s "${EXPORTER_URL}/healthz" | grep -q "ok"; then
        log_success "Exporter 健康检查成功"
    else
        log_error "Exporter 健康检查失败"
        return 1
    fi
    
    # 测试指标端点
    if curl -s "${EXPORTER_URL}/metrics" | grep -q "uwsgi_"; then
        log_success "Exporter 指标端点访问成功"
    else
        log_error "Exporter 指标端点访问失败"
        return 1
    fi
    
    # 检查具体指标
    METRICS=$(curl -s "${EXPORTER_URL}/metrics")
    
    # 检查基础指标
    if echo "${METRICS}" | grep -q "uwsgi_listen_queue"; then
        log_success "基础指标收集正常"
    else
        log_warning "基础指标可能未收集到"
    fi
    
    # 检查 Worker 指标
    if echo "${METRICS}" | grep -q "uwsgi_worker_requests"; then
        log_success "Worker 指标收集正常"
    else
        log_warning "Worker 指标可能未收集到"
    fi
    
    # 检查应用指标
    if echo "${METRICS}" | grep -q "uwsgi_app_requests"; then
        log_success "应用指标收集正常"
    else
        log_warning "应用指标可能未收集到"
    fi
    
    log_success "Exporter 测试完成"
}

# 显示指标
show_metrics() {
    log_info "显示收集的指标..."
    
    echo "=== uWSGI Exporter 指标 ==="
    curl -s "${EXPORTER_URL}/metrics" | grep "uwsgi_" | head -20
    
    echo -e "\n=== 完整指标链接 ==="
    echo "指标页面: ${EXPORTER_URL}/metrics"
    echo "健康检查: ${EXPORTER_URL}/healthz"
    echo "着陆页: ${EXPORTER_URL}/"
}

# 清理函数
cleanup() {
    log_info "清理测试环境..."
    
    # 停止 Exporter
    if pgrep -f "uos_uwsgi_exporter" > /dev/null; then
        pkill -f "uos_uwsgi_exporter"
        log_info "已停止 uWSGI Exporter"
    fi
    
    # 停止 uWSGI
    if pgrep -f "uwsgi.*uwsgi.ini" > /dev/null; then
        pkill -f "uwsgi.*uwsgi.ini"
        log_info "已停止 uWSGI 测试应用"
    fi
    
    # 清理临时文件
    rm -f /tmp/uwsgi_test.sock /tmp/uwsgi_test.pid /tmp/uwsgi_test.log
    rm -f /tmp/uwsgi_daemon.log /tmp/uwsgi_req.log /tmp/uwsgi_error.log
    
    log_success "清理完成"
}

# 主函数
main() {
    log_info "开始 uWSGI Exporter 测试..."
    
    # 设置信号处理
    trap cleanup EXIT INT TERM
    
    # 检查依赖
    check_dependencies
    
    # 安装依赖
    install_python_deps
    
    # 启动 uWSGI
    start_uwsgi
    
    # 测试 uWSGI 应用
    test_uwsgi_app
    
    # 测试统计接口
    test_stats_interface
    
    # 启动 Exporter
    start_exporter
    
    # 测试 Exporter
    test_exporter
    
    # 显示指标
    show_metrics
    
    log_success "所有测试完成！"
    log_info "按 Ctrl+C 停止测试并清理环境"
    
    # 保持运行
    wait
}

# 运行主函数
main "$@" 